package com.example.harmonet.harmtorch;

public class Conv1d implements Layer {
    int length, in_channel, out_channel;
    int pad;
    boolean hasBias;
    int stride;
    float[][] kernel;
    float[] bias;
    int param;


    public Conv1d(int in_channel, int out_channel, int kernel_size, int stride, int padding, boolean bias) {
        this.in_channel = in_channel;
        this.out_channel = out_channel;
        this.length = kernel_size;
        this.stride = stride;
        this.pad = padding;
        this.hasBias = bias;
        this.kernel = new float[in_channel * out_channel][length];
        if (hasBias) {
            this.bias = new float[out_channel];
            this.param = in_channel * out_channel * length + out_channel;
        }
        else {
            this.param = in_channel * out_channel * length;
        }
    }

    private void compute(Tensor in, Tensor out, int n, int i, int l1) throws Exception {
        if(hasBias) {
            out.tensor()[out.pos(n, i, l1)] = bias[i];
        }
        for(int j = 0; j < in_channel; ++j) {
            for(int l2 = 0; l2 < length; ++l2) {
                int l = l1 * stride + l2 - pad;
                if (l >= 0 && l < in.dim()[2]) {
                    out.tensor()[out.pos(n, i, l1)] += in.tensor()[in.pos(n, j, l)]
                            * kernel[i * in_channel + j][l2];
                }
            }
        }
    }

    @Override
    public Tensor forward(Tensor in) throws Exception {
        int[] out_dim = new int[3];
        out_dim[0] = in.dim()[0];
        out_dim[1] = out_channel;
        out_dim[2] = (in.dim()[2] + 2 * pad - length) / stride + 1;
        Tensor out = new Tensor(out_dim);
        for(int n = 0; n < out_dim[0]; ++n) {
            for(int i = 0; i < out_dim[1]; ++i) {
                for(int l1 = 0; l1 < out_dim[2]; ++l1) {
                    compute(in, out, n, i, l1);
                }
            }
        }
        return out;
    }

    @Override
    public int getParam() {
        return param;
    }

    @Override
    public void init(float[] param) {
        int idx = 0;
        for(int i = 0; i < in_channel * out_channel; ++i) {
            for(int j = 0; j < length; ++j) {
                kernel[i][j] = param[idx];
                ++idx;
            }
        }
        for(int i = 0; i < out_channel; ++i) {
            bias[i] = param[idx];
            ++idx;
        }
    }
}
